from model import *
from types import SimpleNamespace
from dataset import *
from torch.utils.data import DataLoader
from utils import *

device='cuda'
batch_size=128
maxlen=50
sample=0
epoch=220

args={}
data="Toys_and_Games"
args['device']=device
args['hidden_units']=50
args['maxlen']=maxlen
args['dropout_rate']=0.2
args['num_blocks']=2
args['num_heads']=1
model_path=f"./re"
args = SimpleNamespace(**args)
dataset = SeqDataset("data/"+data,maxlen=maxlen)
candidate_dataloader=DataLoader(dataset.candidate_item,batch_size=batch_size,collate_fn=collate_fn3)
itemnum=dataset.item_max-1
sasmodel = SASRec(1,itemnum, args).to(args.device)
for name, param in sasmodel.named_parameters():
    try:
        torch.nn.init.xavier_normal_(param.data)
    except:
        pass  # just ignore those failed init layers

sasmodel.pos_emb.weight.data[0, :] = 0
sasmodel.item_emb.weight.data[0, :] = 0
sasmodel.load_state_dict(torch.load(model_path, map_location=torch.device(args.device)))
dataset = data_partition(data)
sasmodel.eval()
t_test = sasrec_evaluate_all(sasmodel, dataset, args)
# t_valid = evaluate_valid(sasmodel, dataset, args)
t_valid=[0,0]

print(f'\nEpoch:{epoch}\ntest (HR@10: {t_test[1]}, nDCG@10: {t_test[0]})')